"""Module for loading simulation data."""

import os
from ..utils import file_utils


class DataLoader:
    """Base class for loading simulation data."""
    
    def __init__(self):
        pass
        
    def load_sky_image(self, file_path, prefix='real'):
        """Load sky image from simulation data.
        
        Args:
            file_path: Path to a FITS file or a folder containing .npy files
            prefix: Prefix for .npy files ('real', 'red', or 'red2'), defaults to 'real'
            
        Returns:
            Data array loaded from the file or folder
        """
        if os.path.isfile(file_path):
            # Check if it's a FITS file
            if file_path.endswith('.fits') or file_path.endswith('.fit'):
                data, header = file_utils.read_fits_file(file_path)
                return data
            else:
                raise ValueError(f"Unsupported file format: {file_path}")
        elif os.path.isdir(file_path):
            # Handle folder with .npy files
            return file_utils.read_npy_file_sequence(file_path, prefix=prefix)
        else:
            raise FileNotFoundError(f"File or directory not found: {file_path}")
        
    def load_spectral_data(self, file_path):
        """Load spectral data from simulation.
        
        Args:
            file_path: Path to filterbank file
            
        Returns:
            Filterbank object
        """
        # For now, only support filterbank files
        # This can be extended to support other spectral data formats
        return file_utils.read_filterbank_file(file_path)